import os
import decimal

import akshare as ak
import pandas as pd
import exchange_calendars as xcals
import requests
import datetime

import draw_display


g_encoding_to_csv = "utf-8-sig"

g_data_base_dir = '.'

g_stock_big_market_value_df : pd.DataFrame


def get_all_stock_info(data_dir):
    sh_a_stocks_df = ak.stock_info_sh_name_code(symbol="主板A股")
    sh_b_stocks_df = ak.stock_info_sh_name_code(symbol="主板B股")
    sh_kc_stocks_df = ak.stock_info_sh_name_code(symbol="科创板")
    sz_a_stocks_df = ak.stock_info_sz_name_code(symbol="A股列表")
    sz_b_stocks_df = ak.stock_info_sz_name_code(symbol="B股列表")
    sz_cdr_stocks_df = ak.stock_info_sz_name_code(symbol="CDR列表")
    sz_ab_stocks_df = ak.stock_info_sz_name_code(symbol="AB股列表")
    bj_stocks_df = ak.stock_info_bj_name_code()

    sh_a_stocks_df.to_csv(data_dir + "sh_主板A股" + ".csv", encoding=g_encoding_to_csv, index=False)
    sh_b_stocks_df.to_csv(data_dir + "sh_主板B股" + ".csv", encoding=g_encoding_to_csv, index=False)
    sh_kc_stocks_df.to_csv(data_dir + "sh_科创板" + ".csv", encoding=g_encoding_to_csv, index=False)
    
    sz_a_stocks_df.to_csv(data_dir + "sz_主板A股" + ".csv", encoding=g_encoding_to_csv, index=False)
    sz_b_stocks_df.to_csv(data_dir + "sz_主板B股" + ".csv", encoding=g_encoding_to_csv, index=False)
    sz_cdr_stocks_df.to_csv(data_dir + "sz_CDR" + ".csv", encoding=g_encoding_to_csv, index=False)
    sz_ab_stocks_df.to_csv(data_dir + "sz_AB股" + ".csv", encoding=g_encoding_to_csv, index=False)
    
    bj_stocks_df.to_csv(data_dir + "bj_股票" + ".csv", encoding=g_encoding_to_csv, index=False)

    zh_a_spot_df = ak.stock_zh_a_spot_em()
    zh_b_spot_df = ak.stock_zh_b_spot_em()
    zh_a_stop_df = ak.stock_zh_a_stop_em()
    zh_a_st_df = ak.stock_zh_a_st_em()
    zh_a_spot_df.to_csv(data_dir + "zh_a_spot" + ".csv", encoding=g_encoding_to_csv, index=False)
    zh_b_spot_df.to_csv(data_dir + "zh_b_spot" + ".csv", encoding=g_encoding_to_csv, index=False)
    zh_a_stop_df.to_csv(data_dir + "zh_a_stop" + ".csv", encoding=g_encoding_to_csv, index=False)
    zh_a_st_df.to_csv(data_dir + "zh_a_st" + ".csv", encoding=g_encoding_to_csv, index=False)

def get_szse_deal_daily(date:str):
    url = "https://www.szse.cn/api/report/ShowReport/data"
    params = {
            "txtQueryDate": "-".join([date[:4], date[4:6], date[6:]]),
            "SHOWTYPE": "JSON",
            "CATALOGID": "scsj_gprdgk_after",
            "tjzqlb": "D",
            "TABKEY": "tab1",
        }
    headers = {
        "Referer": "https://www.szse.cn/market/stock/situation/daily/index.html",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/89.0.4389.90 Safari/537.36",
    }

    r = requests.get(url, params=params, headers=headers)
    data_json = r.json()

    return float(data_json[0]["data"][1]["gp"].replace(",", ""))

def get_stock_deal_volume(start_date:str, end_date:str):
    # 获取交易日历表 XSHG 中国沪市
    xshg = xcals.get_calendar(name="XSHG")
    xshg_range = xshg.schedule.loc[start_date:end_date]
    trade_day_list = xshg_range.index.strftime("%Y%m%d").tolist()
    sse_volume_list = []
    szse_volume_list = []
    total_volume_list = []

    for key in trade_day_list:
        sse_deal_daily_df = ak.stock_sse_deal_daily(date=key)
        sse_volume = sse_deal_daily_df.iloc[3]["股票"]
        szse_volume = get_szse_deal_daily(date=key)
        sse_volume_list.append(sse_volume)
        szse_volume_list.append(szse_volume)
        total_volume_list.append(sse_volume + szse_volume)

    stock_deal_volume_df = pd.DataFrame({
                            "交易日": trade_day_list, 
                            "上交所成交额": sse_volume_list, 
                            "深交所成交额":szse_volume_list, 
                            "两市总成交额":total_volume_list
                         })
    
    stock_deal_volume_df.to_csv(g_data_base_dir + "沪深两市成交额统计.csv", encoding=g_encoding_to_csv)

    return stock_deal_volume_df

def stock_analysis():
    zh_a_spot_filename = g_data_base_dir + "zh_a_spot.csv"
    temp_dict = {}
    first_line = True
    with open(zh_a_spot_filename, "r", encoding="utf8") as zh_a_spot_file:
        for line in zh_a_spot_file:
            line = line.rstrip("\n")
            line_list = line.split(',')
            line_list.pop(0)
            if first_line:
                temp_dict = {key:[] for key in line_list}
                first_line = False
            else:
                for key, value in zip(temp_dict, line_list):
                    if key == "总市值":
                        if value == "":
                            value = 0.0
                        else:
                            value = float(value)/1000e8
                            value = decimal.Decimal(value=value).quantize(decimal.Decimal("0.00"))
                    elif key.find("涨跌幅") != -1:
                        if value == "":
                            value = 0.0
                        else:
                            value = float(value)
                        
                    temp_dict[key].append(value)

    temp_df = pd.DataFrame(temp_dict)
    temp_df.to_csv(g_data_base_dir + "copy_zh_a_spot.csv", encoding=g_encoding_to_csv, index=False)

    global g_stock_big_market_value_df
    g_stock_big_market_value_df = temp_df[temp_df["总市值"] >= 1].sort_values(by="总市值", ascending=False)
    g_stock_big_market_value_df.to_csv(g_data_base_dir + "千亿市值以上.csv", encoding=g_encoding_to_csv, index=False)

    stock_market_value_be_5000_df = temp_df[temp_df["总市值"] >= 5].sort_values(by="总市值", ascending=False)
    stock_market_value_be_1000_df = temp_df[temp_df["总市值"] >= 1].sort_values(by="总市值", ascending=False)
    stock_market_value_be_500_df = temp_df[temp_df["总市值"] >= 0.5].sort_values(by="总市值", ascending=False)
    stock_market_value_be_100_df = temp_df[temp_df["总市值"] >= 0.1].sort_values(by="总市值", ascending=False)
    stock_market_value_be_50_df = temp_df[temp_df["总市值"] >= 0.05].sort_values(by="总市值", ascending=False)
    stock_market_value_lt_50_df = temp_df[temp_df["总市值"] > 0.0].sort_values(by="总市值", ascending=False)
    stock_market_value_bad_df = temp_df[temp_df["总市值"] == 0.0]

    stock_market_value_be_5000_df.to_csv(g_data_base_dir + "stock_market_value_be_5000_df.csv", encoding=g_encoding_to_csv, index=False)
    stock_market_value_be_1000_df.to_csv(g_data_base_dir + "stock_market_value_be_1000_df.csv", encoding=g_encoding_to_csv, index=False)
    stock_market_value_be_500_df.to_csv(g_data_base_dir + "stock_market_value_be_500_df.csv", encoding=g_encoding_to_csv, index=False)
    stock_market_value_be_100_df.to_csv(g_data_base_dir + "stock_market_value_be_100_df.csv", encoding=g_encoding_to_csv, index=False)
    stock_market_value_be_50_df.to_csv(g_data_base_dir + "stock_market_value_be_50_df.csv", encoding=g_encoding_to_csv, index=False)
    stock_market_value_lt_50_df.to_csv(g_data_base_dir + "stock_market_value_lt_50_df.csv", encoding=g_encoding_to_csv, index=False)
    stock_market_value_bad_df.to_csv(g_data_base_dir + "stock_market_value_bad_df.csv", encoding=g_encoding_to_csv, index=False)

    return [stock_market_value_be_5000_df, 
            stock_market_value_be_1000_df,
            stock_market_value_be_500_df,
            stock_market_value_be_100_df,
            stock_market_value_be_50_df,
            stock_market_value_lt_50_df,
            stock_market_value_bad_df]

if __name__ == "__main__":
    print(f'{"Current directory: "} {os.getcwd()}')

    g_data_base_dir = os.getcwd() + "/data/"
    print(os.path.isdir(g_data_base_dir))

    if not os.path.isdir(g_data_base_dir):
        os.mkdir(g_data_base_dir)
        get_all_stock_info(g_data_base_dir)
    
    stock_levels_list = stock_analysis()

    #draw_display.display_market_level(g_data_base_dir, stock_levels_list)
    draw_display.display_price_change(g_data_base_dir, stock_levels_list[1], "60日涨跌幅")
    draw_display.display_price_change(g_data_base_dir, stock_levels_list[1], "年初至今涨跌幅")
    #draw_display.display_big_market_value_stocks(g_data_base_dir, g_stock_big_market_value_df)
    #today_str = datetime.datetime.today().strftime("%Y%m%d")
    #stock_deal_volume_df = get_stock_deal_volume("20240601", today_str)
    #draw_display.display_deal_amount(g_data_base_dir, stock_deal_volume_df)